import pyaudio
import json
import sys
from vosk import Model, KaldiRecognizer, SetLogLevel

# 禁用 Vosk 日志信息（可选）
SetLogLevel(-1)

# 加载 Vosk 语音识别模型（确保使用更大的模型）
MODEL_PATH = "vosk-model-small-cn-0.22"  # 更换为更大的模型路径
model = Model(MODEL_PATH)

# 设置麦克风参数
SAMPLE_RATE = 16000  # Vosk 需要 16kHz 采样率
CHUNK_SIZE = 2048  # 增加性能，减少延迟
FORMAT = pyaudio.paInt16  # 16-bit PCM
CHANNELS = 1  # 单声道输入

# 初始化 Vosk 识别器
recognizer = KaldiRecognizer(model, SAMPLE_RATE)
recognizer.SetWords(True)  # 启用单词输出（可选）

# 初始化 PyAudio
audio = pyaudio.PyAudio()

# 打开麦克风输入流
stream = audio.open(format=FORMAT,
                    channels=CHANNELS,
                    rate=SAMPLE_RATE,
                    input=True,
                    frames_per_buffer=CHUNK_SIZE)

print("🎤 语音识别已启动，开始说话...（按 Ctrl+C 退出）")

try:
    while True:
        # 读取音频数据
        data = stream.read(CHUNK_SIZE, exception_on_overflow=False)

        # 进行语音识别
        partial_result = recognizer.PartialResult()
        partial_result_json = json.loads(partial_result)
        if 'partial' in partial_result_json:
            print("⏳ 中间识别结果:", partial_result_json['partial'])  # 实时打印中间结果

        if recognizer.AcceptWaveform(data):
            result = json.loads(recognizer.Result())
            if 'text' in result:
                print("📝 识别结果:", result['text'])  # 实时打印识别的文本

except KeyboardInterrupt:
    print("\n⏹ 语音识别已停止。")

finally:
    # 关闭音频流和 PyAudio
    stream.stop_stream()
    stream.close()
    audio.terminate()
